#! /usr/bin/env python
# -*- coding: utf-8 -*-

import os
import struct
import numpy as np
from sklearn import svm
import time
from sklearn.model_selection import StratifiedShuffleSplit
from sklearn.model_selection import GridSearchCV


def load_mnist(path, kind='train'):
    '''Load MNIST data set from `path`,
    see http://yann.lecun.com/exdb/mnist/ for file formats'''
    labels_path = os.path.join(path, '%s-labels-idx1-ubyte' % kind)
    images_path = os.path.join(path, '%s-images-idx3-ubyte' % kind)

    with open(labels_path, 'rb') as lbpath:
        _, _ = struct.unpack('>II', lbpath.read(8))
        labels = np.fromfile(lbpath, dtype=np.uint8)

    with open(images_path, 'rb') as imgpath:
        _, num, rows, cols = struct.unpack('>IIII', imgpath.read(16))
        images = np.fromfile(imgpath, dtype=np.uint8).reshape(num, rows * cols)

    return images, labels


def main():
    # Download data sets from http://yann.lecun.com/exdb/mnist/ and uncompress them into ../mnist
    train_images_all, train_labels_all = load_mnist('../mnist', kind='train')
    test_images_all, test_labels_all = load_mnist('../mnist', 't10k')
    # print(train_images_all.shape)
    # print(train_labels_all.shape)
    # print(test_images_all.shape)
    # print(test_labels_all.shape)

    NUM_TRAIN = train_images_all.shape[0]
    NUM_TEST = test_images_all.shape[0]
    train_images = train_images_all[0:NUM_TRAIN, :]
    train_labels = train_labels_all[0:NUM_TRAIN]
    test_images = test_images_all[0:NUM_TEST, :]
    test_labels = test_labels_all[0:NUM_TEST]

    # Binarization
    test_images[test_images > 0] = 1
    train_images[train_images > 0] = 1

    # Find parameters
    # start_time = time.time()
    # print('start: ' + time.strftime('%Y-%m-%d %H:%M:%S'))
    # C_range = 10. ** np.arange(-3, 8)
    # gamma_range = 10. ** np.arange(-5, 4)
    # param_grid = dict(gamma=gamma_range, C=C_range)
    # cv = StratifiedShuffleSplit(n_splits=5, test_size=0.2, random_state=42)
    # grid = GridSearchCV(svm.SVC(), param_grid=param_grid, cv=cv)
    # grid.fit(train_images, train_labels)
    # print("The best parameters are %s with a score of %0.2f" % (grid.best_params_, grid.best_score_))
    # print('end: ' + time.strftime('%Y-%m-%d %H:%M:%S'))
    # print('running: ' + str(time.time() - start_time) + ' seconds')

    start_time = time.time()
    print('start: ' + time.strftime('%Y-%m-%d %H:%M:%S'))
    clf = svm.SVC(C=5, kernel='rbf', gamma=0.05)
    clf.fit(train_images, train_labels)
    print('end: ' + time.strftime('%Y-%m-%d %H:%M:%S'))
    print('(' + str(time.time() - start_time) + ' seconds)')
    print(str(clf.score(test_images, test_labels) * 100.0) + '%')


if __name__ == '__main__':
    main()
